# %%
# 我们会沿用CIFAR-10的教程代码，
# 来实现用tensorboard可视化模型数据和训练过程
# %%
# 导入包库
import matplotlib.pyplot as plt
import numpy as np

import torch
import torchvision
import torchvision.transforms as transforms

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
# %%
# 设置路径
from pathlib import Path

PATH = Path(r'C:\files\git_repository\pytorch-learning\pytorch学习\通过示例学习')
print(PATH)
# %% 数据转换，transforms
# 设置transforms的组合,主要是转化为tensor和进行标准化
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5,), (0.5,))])

# datasets,下载数据集
trainset = torchvision.datasets.FashionMNIST(str(PATH / 'data'),
                                             download=True,
                                             train=True,
                                             transform=transform)
testset = torchvision.datasets.FashionMNIST(str(PATH / 'data'),
                                            download=True,
                                            train=False,
                                            transform=transform)

# dataloaders，设置dataloader
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=0)

testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=0)

# constant for classes，这是10个类别的标签
classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
           'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle Boot')


# %%
# 利用辅助函数来显示一张图片
# helper function to show an image
# (used in the `plot_classes_preds` function below)

def matplotlib_imshow(img, one_channel=False):
    # 对灰度图像和彩色图像分别进行处理
    # one_channel代表是灰度图像
    if one_channel:
        img = img.mean(dim=0)
    # 取消标准化
    img = img / 2 + 0.5  # unnormalize
    # 转换成numpy
    npimg = img.numpy()
    if one_channel:
        plt.imshow(npimg, cmap="Greys")
    else:
        plt.imshow(np.transpose(npimg, (1, 2, 0)))


# %%
# 对模型只进行了小的修改，只是把把原来的现在的图片是一个channel
# 不是3,28x28，也不是32x32
class Net(nn.Module):
    # 定义网络
    def __init__(self):
        super(Net, self).__init__()
        # 实例化层模型
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # 定义前向传播
        # 每层都是卷积接池化
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        # 这里进行了展平
        x = x.view(-1, 16 * 4 * 4)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = Net()
# %%
# 下面定义optimizer和criterion
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# %%
# 下面才是重点
""" 这里才是使用tensorboard的重点 """
# 1.TensorBoard setup 建立
# 用SummaryWriter对象来给TensorBoard写入信息
from torch.utils.tensorboard import SummaryWriter

# 默认的目录是runs，这里我们做的更加具体一点
writer = SummaryWriter(str(PATH / 'runs/fashion_mnist_experiment_1'))
# %%
# 现在让我们给tensorboard写入image，具体来说用make_grid方法
# 取得一些随机的数据

# 使用dataiter，取出一个图像和标签
dataiter = iter(trainloader)
images, labels = next(dataiter)
# 为图像生成网格，就是把多个图像放到一个网格里面
img_grid = torchvision.utils.make_grid(images)
# 显示图像
matplotlib_imshow(img_grid, one_channel=True)

# 写入tensorboard中
writer.add_image('four_fashion_mnist_images', img_grid)

# %%
# in  console
# tensorboard --logdir=runs

# %%
# 使用TensorBoard来检查模型
writer.add_graph(net, images)
writer.close()
#%%
# tensorboard可以很方便的以低维查看高维数据
# 通过add_embedding 方法实现
# 辅助函数
def select_n_random (data,labels,n=100):
    # 从数据集中选择n个随机的数据点和他们相应的标签
    assert len(data) == len(labels)
    perm = torch.randperm(len(data))
    return data[perm][:n], labels[perm][:n]

# %%
